import torch
import cv2
from torchvision.transforms import functional as F

from utils.general import non_max_suppression


def rescale_boxes(boxes, current_dim, original_shape):
    """
    将边界框坐标从当前尺寸缩放回原始图像尺寸
    """
    orig_h, orig_w = original_shape
    current_h, current_w = current_dim
    # 计算缩放比例
    scale = min(orig_w / current_w, orig_h / current_h)
    # 计算偏移量
    dx = (orig_w - scale * current_w) / 2
    dy = (orig_h - scale * current_h) / 2
    # 缩放边界框坐标
    boxes[:, [0, 2]] = (boxes[:, [0, 2]] * scale) + dx
    boxes[:, [1, 3]] = (boxes[:, [1, 3]] * scale) + dy
    return boxes


def split_image(frame, model, augment=False, visualize=False, conf_thres=0.3, iou_thres=0.45, classes=None,
                agnostic_nms=False, max_det=1000):
    """
    对输入的帧进行切图检测

    Args:
        frame (numpy.ndarray): 输入帧，形状为 (H, W, C)。
        model (torch.jit.ScriptModule): 已经导出的 YOLOv5 模型的 TorchScript 格式。
        augment (bool, optional): 是否使用数据增强，默认为 False。
        visualize (bool, optional): 是否可视化结果，默认为 False。
        conf_thres (float, optional): 置信度阈值，默认为 0.3。
        iou_thres (float, optional): IOU 阈值，默认为 0.45。
        classes (list, optional): 检测的类别列表，默认为 None。
        agnostic_nms (bool, optional): 是否使用类别不可知的 NMS，默认为 False。
        max_det (int, optional): 每张图像最大检测框数目，默认为 1000。

    Returns:
        tuple: 切图后的图像列表和检测结果列表，形如 (split_frames, pred_results)。
    """
    # 将帧转换为 Torch 张量
    frame_tensor = F.to_tensor(frame)
    frame_tensor = frame_tensor.unsqueeze(0)
    frame_tensor = frame_tensor.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

    # 模型推理
    with torch.no_grad():
        if augment:
            pred = model(frame_tensor, augment=augment, visualize=visualize)[0]
        else:
            pred = model(frame_tensor)[0]

    # 非极大值抑制
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

    # 处理切图检测结果
    split_frames = []
    pred_results = []
    if pred is not None and len(pred) > 0:
        for i in range(len(pred)):
            if visualize:
                pred[i] = rescale_boxes(pred[i], frame_tensor.shape[-2:], frame.shape[:2])
            for *xyxy, conf, cls in pred[i]:
                x, y, x_, y_ = map(int, xyxy)
                split_frame = frame[y:y_, x:x_]
                split_frames.append(split_frame)
                pred_results.append((x, y, x_, y_, conf, cls))

    return split_frames, pred_results
